#include <zenopy/UsdAPI.hpp>
#include <zenox/minilog.hpp>
#include <zenopy/PythonClient.hpp>
#include <cppcodec/base64_rfc4648.hpp>
#include <zenox/glm.hpp>
#include <zenox/ziputils.hpp>
#include <cstring>
#include <vector>
#include <string>

namespace zenopy {

#define ZENO_PYCLI(...) __VA_ARGS__

std::vector<std::string> stagePrimList(PythonClient &cli, std::string const &stageFile) {
    Json::Value args;
    args["stageFile"] = stageFile;
    args = cli.exec(ZENO_PYCLI(R"PYTHON(
        stage = cache.usdStageOpen(args["stageFile"])
        prims = []
        for prim in stage.Traverse():
            prims.append(str(prim.GetPath()))
        return {"primPaths": prims}
)PYTHON"), args);
    std::vector<std::string> ret;
    ret.reserve(args["primPaths"].size());
    for (auto const &prim: args["primPaths"]) {
        ret.push_back(prim.asString());
    }
    return ret;
}

std::string stageDefaultPrim(PythonClient &cli, std::string const &stageFile) {
    Json::Value args;
    args["stageFile"] = stageFile;
    args = cli.exec(ZENO_PYCLI(R"PYTHON(
        stage = cache.usdStageOpen(args["stageFile"])
        prim = stage.GetDefaultPrim()
        primPath = str(prim.GetPath())
        return {"primPath": primPath}
)PYTHON"), args);
    return args["primPath"].asString();
}

std::string stagePrimName(PythonClient &cli, std::string const &stageFile, std::string const &primPath) {
    Json::Value args;
    args["stageFile"] = stageFile;
    args["primPath"] = primPath;
    args = cli.exec(ZENO_PYCLI(R"PYTHON(
        stage = cache.usdStageOpen(args["stageFile"])
        primPath = args["primPath"]
        prim = stage.GetPrimAtPath(primPath)
        primName = prim.GetName()
        return {"primName": primName}
)PYTHON"), args);
    return args["primName"].asString();
}

std::string stagePrimPropType(PythonClient &cli, std::string const &stageFile, std::string const &primPath, std::string const &propName) {
    Json::Value args;
    args["stageFile"] = stageFile;
    args["primPath"] = primPath;
    args["propName"] = propName;
    args = cli.exec(ZENO_PYCLI(R"PYTHON(
        stage = cache.usdStageOpen(args["stageFile"])
        primPath = args["primPath"]
        propName = args["propName"]
        prim = stage.GetPrimAtPath(primPath)
        if not prim.HasProperty(propName):
            return {"propType": ""}
        prop = prim.GetProperty(propName)
        propType = str(prop.GetTypeName())
        return {"propType": propType}
)PYTHON"), args);
    return args["propType"].asString();
}

std::string stagePrimPropValString(PythonClient &cli, std::string const &stageFile, std::string const &primPath, std::string const &propName) {
    zenox::log_trace("stagePrimPropValString invoking python server with [{}] [{}] [{}]", stageFile, primPath, propName);
    Json::Value args;
    args["stageFile"] = stageFile;
    args["primPath"] = primPath;
    args["propName"] = propName;
    args = cli.exec(ZENO_PYCLI(R"PYTHON(
        stage = cache.usdStageOpen(args["stageFile"])
        primPath = args["primPath"]
        propName = args["propName"]
        prim = stage.GetPrimAtPath(primPath)
        if not prim.HasProperty(propName):
            return {"propValString": ""}
        prop = prim.GetProperty(propName)
        propValString = str(prop.Get())
        return {"propValString": propValString}
)PYTHON"), args);
    return args["propValString"].asString();
}

std::vector<std::string> stagePrimPropValStrList(PythonClient &cli, std::string const &stageFile, std::string const &primPath, std::string const &propName) {
    zenox::log_trace("stagePrimPropValStrList invoking python server with [{}] [{}] [{}]", stageFile, primPath, propName);
    Json::Value args;
    args["stageFile"] = stageFile;
    args["primPath"] = primPath;
    args["propName"] = propName;
    args = cli.exec(ZENO_PYCLI(R"PYTHON(
        stage = cache.usdStageOpen(args["stageFile"])
        primPath = args["primPath"]
        propName = args["propName"]
        prim = stage.GetPrimAtPath(primPath)
        if not prim.HasProperty(propName):
            return {"propValStrList": []}
        prop = prim.GetProperty(propName)
        propVal = prop.Get()
        if propVal is None:
            return {"propValStrList": []}
        propValStrList = list(map(str, list(propVal)))
        return {"propValStrList": propValStrList}
)PYTHON"), args);
    std::vector<std::string> items;
    items.reserve(args["propValStrList"].size());
    for (auto const &item: args["propValStrList"]) {
        items.push_back(item.asString());
    }
    return items;
}

std::vector<size_t> details::_impl_stagePrimPropValArray(PythonClient &cli, std::string const &stageFile, std::string const &primPath, std::string const &propName, std::function<void *(size_t)> const &outDelegate, int assumeShapeDim, int glmComponents, int Sca0Vec1Mat2, int glmCols, int glmRows, size_t sizeofT0) {
    zenox::log_trace("stagePrimPropValArray invoking python server with [{}] [{}] [{}]", stageFile, primPath, propName);
    Json::Value args;
    args["stageFile"] = stageFile;
    args["primPath"] = primPath;
    args["propName"] = propName;
    args = cli.exec(ZENO_PYCLI(R"PYTHON(
        stage = cache.usdStageOpen(args["stageFile"])
        primPath = args["primPath"]
        propName = args["propName"]
        prim = stage.GetPrimAtPath(primPath)
        if not prim.HasProperty(propName):
            return {"propValType": ""}
        prop = prim.GetProperty(propName)
        propValArray = np.array(prop.Get())
        propValType = np.dtype(propValArray.dtype)
        propValItemSize = int(propValType.itemsize)
        propValArraySize = int(propValArray.size)
        propValArrayShape = list(propValArray.shape)
        propValType = str(propValType.name)
        propValArray = base64.b64encode(zlib.compress(propValArray.tobytes(), level=1)).decode()
        return {"propValType": propValType, "propValItemSize": propValItemSize, "propValArraySize": propValArraySize, "propValArrayShape": propValArrayShape, "propValArray": propValArray}
)PYTHON"), args);
    auto valType = args["propValType"].asString();
    if (valType.empty()) [[unlikely]] {
        zenox::log_trace("got empty valType for {}, assuming property not exist and return empty vector", propName);
        std::vector<size_t> shape(assumeShapeDim < 0 ? 0 : assumeShapeDim, 0);
        return shape;
    }
    auto valArr = args["propValArray"].asString();
    auto valArrShape = args["propValArrayShape"];
    auto valArrSize = args["propValArraySize"].asUInt64();
    auto valItemSize = args["propValItemSize"].asUInt64();
    if (sizeofT0 != valItemSize) [[unlikely]] zenox::fatal_error("[{}] item size mismatch: {} != {}", primPath + '@' + propName, sizeofT0, valItemSize);
    std::vector<size_t> shape;
    shape.reserve(valArrShape.size());
    for (auto const &dim: valArrShape) {
        shape.push_back(dim.asUInt64());
    }

    if (Sca0Vec1Mat2 == 1) {
        if (shape.empty()) [[unlikely]] zenox::fatal_error("[{}] shape should not be empty when expecting vec as element type", primPath + '@' + propName);
        if (shape.back() != (size_t)glmComponents) [[unlikely]] zenox::fatal_error("[{}] last dimention of shape mismatch: {} != {}", primPath + '@' + propName, shape.back(), glmComponents);
        shape.pop_back();
    } else if (Sca0Vec1Mat2 == 2) {
        if (shape.size() < 2) [[unlikely]] zenox::fatal_error("[{}] shape should have at least 2-dims when expecting vec as element type", primPath + '@' + propName);
        if (shape.back() != (size_t)glmRows) [[unlikely]] zenox::fatal_error("[{}] last dimention of shape mismatch: {} != {}", primPath + '@' + propName, shape.back(), glmRows);
        shape.pop_back();
        if (shape.back() != (size_t)glmCols) [[unlikely]] zenox::fatal_error("[{}] secondary last dimention of shape mismatch: {} != {}", primPath + '@' + propName, shape.back(), glmCols);
        shape.pop_back();
    }
    if (assumeShapeDim != -1 && shape.size() != (size_t)assumeShapeDim) [[unlikely]] zenox::fatal_error("[{}] shape dimension mismatch: {} != {}", primPath + '@' + propName, shape.size(), assumeShapeDim);

    zenox::log_trace("base64 encoded size {}", valArr.size());
    auto zipArr = cppcodec::base64_rfc4648::decode(valArr);
    zenox::log_trace("base64 decoded size {}", zipArr.size());
    std::vector<uint8_t> rawArr(valArrSize * valItemSize);
    size_t outElems = valArrSize / glmComponents;
    auto *outPtr = (Bytef *)outDelegate(outElems);
    uLongf outSize = outElems * (valItemSize * glmComponents);
    if (valArrSize * valItemSize != outSize) [[unlikely]] zenox::fatal_error("[{}] pre size mismatch: {} != {}", primPath + '@' + propName, valArrSize * valItemSize, outSize);
    if (int status = ::uncompress(outPtr, &outSize, zipArr.data(), zipArr.size()); status != Z_OK) {
        zenox::fatal_error("[{}] zlib uncompress failed, code {}", primPath + '@' + propName, zenox::zliberror(status));
    }
    if (valArrSize * valItemSize != outSize) [[unlikely]] zenox::fatal_error("[{}] post size mismatch: {} != {}", primPath + '@' + propName, valArrSize * valItemSize, outSize);
    zenox::log_trace("zlib uncompressed size {}", outSize);
    return shape;
}

std::vector<std::string> stagePrimPropNames(PythonClient &cli, std::string const &stageFile, std::string const &primPath) {
    Json::Value args;
    args["stageFile"] = stageFile;
    args["primPath"] = primPath;
    args = cli.exec(ZENO_PYCLI(R"PYTHON(
        stage = cache.usdStageOpen(args["stageFile"])
        primPath = args["primPath"]
        prim = stage.GetPrimAtPath(primPath)
        children = []
        for propName in prim.GetPropertyNames():
            children.append(propName)
        return {"primPropNames": children}
)PYTHON"), args);
    std::vector<std::string> ret;
    ret.reserve(args["primPropNames"].size());
    for (auto const &prim: args["primPropNames"]) {
        ret.push_back(prim.asString());
    }
    return ret;
}

std::vector<std::string> stagePrimChildren(PythonClient &cli, std::string const &stageFile, std::string const &primPath) {
    Json::Value args;
    args["stageFile"] = stageFile;
    args["primPath"] = primPath;
    args = cli.exec(ZENO_PYCLI(R"PYTHON(
        stage = cache.usdStageOpen(args["stageFile"])
        primPath = args["primPath"]
        prim = stage.GetPrimAtPath(primPath)
        children = []
        for child in prim.GetChildren():
            children.append(str(child.GetPath()))
        return {"primChildPaths": children}
)PYTHON"), args);
    std::vector<std::string> ret;
    ret.reserve(args["primChildPaths"].size());
    for (auto const &prim: args["primChildPaths"]) {
        ret.push_back(prim.asString());
    }
    return ret;
}

std::string stagePrimChild(PythonClient &cli, std::string const &stageFile, std::string const &primPath, std::string const &childRelName) {
    // return value can be "Xform" "Mesh" etc.
    Json::Value args;
    args["stageFile"] = stageFile;
    args["primPath"] = primPath;
    args["childRelName"] = childRelName;
    args = cli.exec(ZENO_PYCLI(R"PYTHON(
        stage = cache.usdStageOpen(args["stageFile"])
        primPath = args["primPath"]
        childRelName = args["childRelName"]
        prim = stage.GetPrimAtPath(primPath)
        return {"primChild": str(prim.GetChild(childRelName).GetPath())}
)PYTHON"), args);
    return args["primChild"].asString();
}

std::string stagePrimParent(PythonClient &cli, std::string const &stageFile, std::string const &primPath) {
    // return value can be "Xform" "Mesh" etc.
    Json::Value args;
    args["stageFile"] = stageFile;
    args["primPath"] = primPath;
    args = cli.exec(ZENO_PYCLI(R"PYTHON(
        stage = cache.usdStageOpen(args["stageFile"])
        primPath = args["primPath"]
        prim = stage.GetPrimAtPath(primPath)
        return {"primParent": str(prim.GetParent().GetPath())}
)PYTHON"), args);
    return args["primParent"].asString();
}

std::string stagePrimType(PythonClient &cli, std::string const &stageFile, std::string const &primPath) {
    Json::Value args;
    args["stageFile"] = stageFile;
    args["primPath"] = primPath;
    args = cli.exec(ZENO_PYCLI(R"PYTHON(
        stage = cache.usdStageOpen(args["stageFile"])
        primPath = args["primPath"]
        prim = stage.GetPrimAtPath(primPath)
        return {"primType": prim.GetTypeName()}
)PYTHON"), args);
    return args["primType"].asString();
}

bool stagePrimIsFlag(PythonClient &cli, std::string const &stageFile, std::string const &primPath, std::string const &flagType) {
    Json::Value args;
    args["stageFile"] = stageFile;
    args["primPath"] = primPath;
    args["flagType"] = flagType;
    args = cli.exec(ZENO_PYCLI(R"PYTHON(
        stage = cache.usdStageOpen(args["stageFile"])
        primPath = args["primPath"]
        flagType = args["flagType"]
        prim = stage.GetPrimAtPath(primPath)
        return {"primIsFlag": getattr(prim, "Is" + flagType)()}
)PYTHON"), args);
    return args["primIsFlag"].asBool();
}

}
